#!/bin/bash
set -e

# ============================================================
# HunyuanWorld-Mirror-Ascend NPU 部署脚本
# 对外暴露端口: 18003
# 启动模式: api (REST API) 或 gradio (Web UI)
# ============================================================

# 配置
IMAGE_NAME="hunyuan-world-mirror-ascend:latest"
CONTAINER_NAME="hunyuan-world-mirror-api"
CODE_PATH=$(pwd)/HunyuanWorld-Mirror # 假设代码在当前目录的子目录中
HOST_PORT=18003
CONTAINER_PORT=8000  # FastAPI 默认端口

# 模型路径（需要提前下载模型到此目录）
MODEL_HOST_PATH="/root/workspace/HunyuanWorld-Mirror"

# 启动模式: api 或 gradio
START_MODE=${1:-"api"}

# 1. 检查/下载代码
if [ ! -d "$CODE_PATH" ]; then
    echo "Cloning HunyuanWorld-Mirror..."
    git clone https://github.com/Tencent-Hunyuan/HunyuanWorld-Mirror.git "$CODE_PATH" || echo "Clone failed, will try inside container or check manually."
fi

# 2. 构建镜像
echo "Building Docker image..."
docker build -t "$IMAGE_NAME" -f HunyuanWorld-Mirror-Ascend/Dockerfile HunyuanWorld-Mirror-Ascend/

# 3. 运行容器
echo "Running container..."
# 清理旧容器
docker stop "$CONTAINER_NAME" 2>/dev/null || true
docker rm "$CONTAINER_NAME" 2>/dev/null || true

# 指定使用 NPU 0（根据实际可用设备调整）
NPU_ID=1
NPU_DEVICE_ARGS="--device=/dev/davinci${NPU_ID}"
if [ -e "/dev/davinci_manager" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/davinci_manager"
fi
if [ -e "/dev/devmm_svm" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/devmm_svm"
fi
if [ -e "/dev/hisi_hdc" ]; then
    NPU_DEVICE_ARGS="$NPU_DEVICE_ARGS --device=/dev/hisi_hdc"
fi

docker run -d \
  --name "$CONTAINER_NAME" \
  --restart=unless-stopped \
  --privileged \
  -p ${HOST_PORT}:${CONTAINER_PORT} \
  $NPU_DEVICE_ARGS \
  -e ASCEND_VISIBLE_DEVICES=$NPU_ID \
  -e ASCEND_RT_VISIBLE_DEVICES=$NPU_ID \
  -e ASCEND_DEVICE_ID=$NPU_ID \
  -e START_MODE=$START_MODE \
  -e MODEL_PATH=/app/models/HunyuanWorld-Mirror \
  -e USE_FP16=true \
  -v /usr/local/dcmi:/usr/local/dcmi \
  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
  -v /etc/ascend_install.info:/etc/ascend_install.info \
  -v "$CODE_PATH":/app/HunyuanWorld-Mirror \
  -v "$MODEL_HOST_PATH":/app/models/HunyuanWorld-Mirror \
  "$IMAGE_NAME"

echo ""
echo "============================================================"
echo "Container $CONTAINER_NAME started in background."
echo "Mode: $START_MODE"
echo ""
echo "API Endpoints:"
echo "  健康检查:    http://localhost:${HOST_PORT}/health"
echo "  API文档:     http://localhost:${HOST_PORT}/docs"
echo "  3D重建:      http://localhost:${HOST_PORT}/v1/3d/reconstruct"
echo "  深度估计:    http://localhost:${HOST_PORT}/v1/3d/depth"
echo "  法线估计:    http://localhost:${HOST_PORT}/v1/3d/normal"
echo "============================================================"
echo ""
echo "查看日志:     docker logs -f $CONTAINER_NAME"
echo "停止容器:     docker stop $CONTAINER_NAME"
echo "重启容器:     docker restart $CONTAINER_NAME"
echo ""
echo "最近日志 (最后20行):"
sleep 2
docker logs --tail 20 "$CONTAINER_NAME"





